# coding: utf-8

from __future__ import print_function
from mll import *
from dataset import *
from logit import *
from matplotlib import pyplot as plt
import pickle

X = np.insert(X, 0, 1, 1)

def test_logit():
    train_logit()
    plot_logit()

def train_logit():
    # alpha=0.001, max_iter=500000
    # loss=0.431984898504
    clf = Logistic(alpha=0.001, max_iter=500000, debug=0)

    clf.fit(X, Y)
    
    with open('logit.pkl', 'wb') as f:
        pickle.dump(clf, f)
        
    print('done..')


def plot_logit():
    
    with open('logit.pkl', 'rb') as f:
        clf = pickle.load(f, fix_imports=True)
    
    Y_test = clf.predict(X)
    print('acc:', (Y_test == Y).mean())
    # 0.8125
    
    clr_test = [('r' if yi == 1 else 'b') for yi in Y]
    
    plt.scatter(X.T[1], X.T[2], c=clr_test) 
    
    x_min = np.min(X.T[1])
    x_max = np.max(X.T[1])
    decision_x = np.asarray([x_min, x_max])
    
    theta0, theta1, theta2 = clf.theta.ravel()
    decision_y = - (theta0 + theta1 * decision_x) / theta2
    plt.plot(decision_x, decision_y, 'g')
    plt.show()
    

def test_mll():
    train_mll()
    plot_mll()

def train_mll():
    # alpha=0.001, max_iter=1000000
    # loss=0.420520277699
    clf = MultiLayerLogistic(alpha=0.001, max_iter=1000000, debug=False)
    
    clf.fit(X, Y)
    
    with open('mll.pkl', 'wb') as f:
        pickle.dump(clf, f)
        
    print('done..')

    
def plot_mll():
    
    with open('mll.pkl', 'rb') as f:
        clf = pickle.load(f)
    
    Y_test = clf.predict(X)
    # print(Y_test)
    print('acc:', (Y_test == Y).mean())
    # 0.825
    
    clr_test = [('r' if l == 1 else 'b') for l in Y]
    
    plt.scatter(X.T[1], X.T[2], c=clr_test)
        
    x_min = np.min(X.T[1])
    x_max = np.max(X.T[1]) + 1
    y_min = np.min(X.T[2])
    y_max = np.max(X.T[2]) + 1
    
    x_rng = np.arange(x_min, x_max)
    y_rng = np.arange(y_min, y_max)
    x_rng, y_rng = np.meshgrid(x_rng, y_rng)
    x_rng, y_rng = x_rng.ravel(), y_rng.ravel()
    X_pred = np.asarray([np.ones(x_rng.shape), x_rng, y_rng]).T
    
    c_rng = clf.predict(X_pred)
    c_rng = [('r' if l == 1 else 'b') for l in c_rng]
    
    plt.scatter(x_rng, y_rng, c=c_rng, alpha=0.25)
    
    plt.show()

plot_mll()

